from typing import TypedDict
from IPython.display import Image,display
from langgraph.constants import END, START
from langgraph.graph.state import CompiledStateGraph, StateGraph


#1 。  定义状态
class ElephantInFridgeState(TypedDict):
    """
    定义状态--  大象在冰箱中
    状态内的参数可以在多个节点保存和流转
    """
    fridge_open:bool  #冰箱门是否打开
    elphant_inside:bool #大象是否在冰箱

# 2. 定义每个步骤的处理节点
def open_fridge(state: ElephantInFridgeState) -> ElephantInFridgeState:
    """
    第一步：把冰箱门打开
    """
    print('正在打开冰箱...')
    print("冰箱门已打开")
    state["fridge_open"] = True
    return state

def put_elepant(state: ElephantInFridgeState) -> ElephantInFridgeState:
    """第二布：把大象放入冰箱"""
    print("正在把大象放入冰箱")
    print("已经把大象放入冰箱")
    state["elphant_inside"] = True
    return state

def close_fridge(state: ElephantInFridgeState)-> ElephantInFridgeState:
    print('正在关冰箱。。。')
    print('已经关好冰箱')
    state["elphant_inside"] = True

#3.  构建graph工作流图
def build_graph() -> CompiledStateGraph:
    graph = StateGraph(ElephantInFridgeState)
    #添加节点
    graph.add_node("open_fridge",open_fridge)
    graph.add_node("put_elepant",put_elepant)
    graph.add_node("close_fridge",close_fridge)
    #添加边
    graph.add_edge(START, "open_fridge")
    graph.add_edge("open_fridge", "put_elepant")
    graph.add_edge("put_elepant", "close_fridge")
    graph.add_edge("close_fridge", END)
    return graph.compile()

#4. 运行工作流
def run_workflow():
    """运行工作流"""
    #构建图
    graph = build_graph()
    #定义初始状体
    initial_state = ElephantInFridgeState(fridge_open=False,elphant_inside=False)
    #运行图
    result = graph.invoke(initial_state)
    # 输出状态
    print('最终状体\n')
    png_data = graph.get_graph().draw_mermaid_png()
    with open('jack.png', "wb") as f:
        f.write(png_data)
    return ""

if __name__ == "__main__":
    run_workflow()